/*
    DLL Injector[Rewrite Version] in Rust.
    @5mukx
*/

use std::env::args;
use std::ffi::CString;
use std::ptr::null_mut;
use winapi::ctypes::c_void;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::CloseHandle;
use winapi::um::libloaderapi::{GetModuleHandleA, GetProcAddress};
use winapi::um::memoryapi::{VirtualAllocEx, VirtualFreeEx, WriteProcessMemory};
use winapi::um::processthreadsapi::{
    CreateRemoteThread, GetExitCodeThread, OpenProcess, TerminateThread,
};
use winapi::um::synchapi::WaitForSingleObject;
use winapi::um::winbase::INFINITE;
use winapi::um::winnt::{MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_READWRITE, PROCESS_ALL_ACCESS};

macro_rules! log {
    ($level:expr, $msg:expr) => {{
        let color = match $level {
            "DEBUG" => "\x1b[94m",
            "INFO" => "\x1b[32m",
            "WARN" => "\x1b[33m",
            "ERROR" => "\x1b[31m",
            _ => "\x1b[0m",
        };

        match $level {
            "DEBUG" => println!("{}[||] {}\x1b[0m", color, $msg),
            "INFO" => println!("{}[+] {}\x1b[0m", color, $msg),
            "WARN" => println!("{}[!] {}\x1b[0m", color, $msg),
            "ERROR" => println!("{}[-] {}\x1b[0m", color, $msg),
            _ => println!("{}{}\x1b[0m", color, $msg),
        };
    }};
}

fn main() {
    let args: Vec<String> = args().collect();
    if args.len() != 3 {
        log!("ERROR", "Usage: dll_inject.exe <PID> <DLL Path>");
        return;
    }

    let pid: u32 = match args[1].parse() {
        Ok(id) => id,
        Err(_) => {
            log!("ERROR", "Failed to parse PID");
            return;
        }
    };

    let dll_path = match CString::new(&*args[2]) {
        Ok(path) => path,
        Err(_) => {
            log!("ERROR", "Failed to convert DLL path to CString");
            return;
        }
    };

    let dll_path = dll_path.to_str().expect("Error Converting CString to Str");
    let dllsize = dll_path.len();

    log!("INFO", format!("PID: {}", pid));
    log!("INFO", format!("DLL Path: {}", dll_path));

    unsafe {
        let process: *mut c_void = OpenProcess(PROCESS_ALL_ACCESS, 0, pid);

        if process == NULL {
            log!(
                "ERROR",
                format!(
                    "Failed to get handle to the process: Error Code {}",
                    GetLastError()
                )
            );
            return;
        }

        log!("INFO", format!("HANDLE of {}: {:?}", pid, process));

        let buffer = VirtualAllocEx(
            process,
            null_mut(),
            dllsize,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE,
        );

        if buffer == null_mut() {
            log!(
                "ERROR",
                format!("Failed to allocate buffer: Error Code {}", GetLastError())
            );
            CloseHandle(process);
            return;
        }

        log!("INFO", format!("Buffer Allocated: {:#?}", buffer));

        let write_process = WriteProcessMemory(
            process,
            buffer,
            dll_path.as_ptr() as *const c_void,
            dllsize,
            null_mut(),
        );

        if write_process == 0 {
            log!(
                "ERROR",
                format!(
                    "Failed to write DLL path to process memory: Error Code {}",
                    GetLastError()
                )
            );
            VirtualFreeEx(process, buffer, 0, MEM_RELEASE);
            CloseHandle(process);
            return;
        }

        log!("INFO", format!("Wrote {} to process memory", dll_path));

        let kernel32 = GetModuleHandleA("kernel32.dll\0".as_ptr() as *const _);

        if kernel32.is_null() {
            log!(
                "ERROR",
                format!(
                    "Failed to get handle to kernel32.dll: Error Code {}",
                    GetLastError()
                )
            );
            VirtualFreeEx(process, buffer, 0, MEM_RELEASE);
            CloseHandle(process);
            return;
        }

        log!(
            "INFO",
            format!("Got handle to kernel32.dll: {:#?}", kernel32)
        );

        let load_library_addr = GetProcAddress(kernel32, "LoadLibraryA\0".as_ptr() as *const _);

        if load_library_addr.is_null() {
            log!(
                "ERROR",
                format!(
                    "Failed to get address of LoadLibraryA: Error Code {}",
                    GetLastError()
                )
            );
            VirtualFreeEx(process, buffer, 0, MEM_RELEASE);
            CloseHandle(process);
            return;
        }

        log!(
            "INFO",
            format!("LoadLibraryA Address: {:#?}", load_library_addr)
        );

        let thread = CreateRemoteThread(
            process,
            null_mut(),
            0,
            Some(std::mem::transmute(load_library_addr)),
            buffer,
            0,
            null_mut(),
        );

        if thread.is_null() {
            log!(
                "ERROR",
                format!(
                    "Failed to create remote thread: Error Code {}",
                    GetLastError()
                )
            );
            VirtualFreeEx(process, buffer, 0, MEM_RELEASE);
            CloseHandle(process);
            return;
        }

        log!("INFO", format!("Remote thread created: {:#?}", thread));

        WaitForSingleObject(thread, INFINITE);

        let mut exit_code = 0;

        GetExitCodeThread(thread, &mut exit_code);

        if exit_code == 0x00000103 {
            log!("INFO", "Thread is still active, terminating it.");
            TerminateThread(thread, 0);
        }

        VirtualFreeEx(process, buffer, 0, MEM_RELEASE);
        CloseHandle(thread);
        CloseHandle(process);

        log!("INFO", "DLL Injection executed successfully!");
        return;
    }
}
